# -*- coding: utf-8 -*-
import asyncio
from datetime import datetime

from checker import Checker

import wqb

import settings_config

from generator import Generator
from exports import ExportFiles
import utils
from submitter import Submitter
from improvement import Improvement

def main():

    try:
        print("🚀 启动 WorldQuant Brain 程序")

        credentials = str(input("\n请输入认证文件路径(默认: ~/.brain_credentials.txt)"))
        
        if credentials == "":
            credentials = "~/.brain_credentials.txt"

        wqbs= wqb.WQBSession((utils.load_credentials(credentials)), logger=wqb.wqb_logger(name='logs/wqb_' + datetime.now().strftime('%Y%m%d')))

        print("\n📋 请选择运行模式:")
        print("1: 生成Alpha")
        print("2: 自相关性检查")
        print("3: 自动提交")
        print("4: 生成数据集文件")
        print("5: 导出已提交的Alpha")

        mode = int(input("\n请选择模式 (1-5): "))
        if mode not in [1, 2, 3, 4,5]:
            print("❌ 无效的模式选择")
            return

        if mode == 1:
            
            print(f"\n📋 请选择生成类别：")
            print("1: 来自文件")
            print("2: 一阶")
            print("3: 二阶剪枝")
            print("4: 三阶剪枝")
            gen_mode = int(input("\n请选生成类别 (1-4): "))
            if gen_mode not in [1, 2, 3, 4]:
                print("❌ 无效的生成类别")
                return
            if gen_mode == 1 or gen_mode == 2:
                print("\n📊 可用setting列表:")
                for dataset in settings_config.get_all_settings():
                    print(dataset)
                setting_index = int(input("\n请选择setting编号: "))
                setting = settings_config.get_by_index(setting_index)
                if setting is None:
                    print("❌ 无效的setting编号")
                    return
                generator = Generator(wqbs)
                if gen_mode == 2:
                    dataset_id = input("\n请输入数据集ID: ")
                    if not dataset_id:
                        print("❌ 无效的数据集编号")
                        return
                    
                    list=generator.generate_first(dataset_id, setting)
                else:
                    file_path = str(input("\n请输入文件路径: "))
                    fields = []
                    with open(file_path, "r") as f:
                        for line in f.readlines():
                            fields.append(line.strip())
                    list=generator.generate_first_with_fields(fields, setting)
            else:
                today = datetime.strftime(datetime.now(), "%Y-%m-%d")
                begen_date = input("\n请输入开始日期(YYYY-MM-DD): ")
                if begen_date == "":
                    begen_date = today
                end_date = input("\n请输入结束日期(YYYY-MM-DD): ")
                if end_date == "":
                    end_date = today
                limit_str = input("\n请输入数据量(默认: 1000): ")
                limit = 1000
                if limit_str != '':
                    limit = int(limit_str)

                improvement = Improvement(
                    wqbs
                    , begin_time=datetime.fromisoformat(f'{begen_date}T00:00:00-05:00')
                    , end_time=datetime.fromisoformat(f'{end_date}T23:59:59-05:00')
                    ,limit=limit
                )
                if gen_mode == 3:
                    list=improvement.first_improve()
                else:
                    list=improvement.second_improve()
            # print(f"\n📋 生成Alpha数量: {len(list)}")
            # if len(list) == 0:
            #     return
            # concurrency = int(input("\n📋 请输入回测并发数: "))
            # if concurrency > 8:
            #     print(f"❌ 并发数不能超过8")
            #     concurrency = 8
            # simulator = Simulator(wqbs, concurrency)
            # asyncio.run(simulator.simulate(list))
            
        elif mode == 2:
            print(f"\n📋 请选择检查模式：")
            print("1: 本地检查")
            print("2: 服务器检查")
            check_mode = int(input("\n请选择检查模式 (1-2): "))
            if check_mode not in [1, 2]:
                print("❌ 无效的检查模式")
                return
            # checker = Checker(wqbs,'./results', True)
            # checker.check(check_mode)
        elif mode == 3:

            today = datetime.strftime(datetime.now(), "%Y-%m-%d")
            sharpe = 1.58
            sharpe_str = input("\n请输入Sharpe阈值(默认: 1.58): ")
            if sharpe_str != "":
                sharpe = float(sharpe_str)
            fitness = 1.0
            fitness_str = input("\n请输入Fitness阈值(默认: 1.0): ")
            if fitness_str != "":
                fitness = float(fitness_str)
            self_corr = 0.6
            self_corr_str = input("\n请输入SelfCorr阈值(默认: 0.6): ")
            if self_corr_str != "":
                self_corr = float(self_corr_str)
            print(f"📋 输入的阈值: sharpe={sharpe}, fitness={fitness}, self_corr={self_corr}")
            begen_date = input("\n请输入开始日期(YYYY-MM-DD): ")
            if begen_date == "":
                begen_date = today
            end_date = input("\n请输入结束日期(YYYY-MM-DD): ")
            if end_date == "":
                end_date = today
            Submitter(
                wqbs, 
                begin_time=f"{begen_date}T00:00:00-05:00",
                end_time=f"{end_date}T23:59:59-05:00"
                , sharpe=sharpe
                , fitness=fitness
                , self_corr_threshold=self_corr
            ).submit()
        else:
                
            # 生成数据集文件
            out_put_path = str(input("\n请输入保存文件路径(默认: ./datasetFile): "))
            if out_put_path == "":
                out_put_path = "./datasetFile"
            export = ExportFiles(
                wqbs
                , out_put_path=out_put_path
            )
            if mode == 4:
                export.export_submitted_alphas()
            else:
                export.generate()

    except Exception as e:
        print(f"❌ 程序运行出错: {str(e)}")

if __name__ == '__main__':
    main()